#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_sort_long.h>

#include "halofinder.h"

#ifndef HAVEFITSIO
static char formatString[] = "%ld\t%.10e\t%.10e\t%.10e\t%.10e\t%.10e\t%.10e\t%.10e\t%.10e\t%.10e\n";
static char formatStringRead[] = "%ld\t%e\t%e\t%e\t%e\t%e\t%e\t%e\t%e\t%e\n";
#endif

static void readpeakshalos(char name[]);

/*gets the number of lines in a file*/
long fnumlines(FILE *fp)
{
  long i=-1;
  char c[5000];
  while(!feof(fp))
    {
      ++i;
      fgets(c,5000,fp);
    }
  rewind(fp);
  return i;
}

#ifndef HAVEFITSIO
static void writepeakshalos_ascii(char name[])
{
  FILE *fp;
  int currTask;
  long i;
  
  if(ThisTask == 0)
    remove(name);
  
  MPI_Barrier(MPI_COMM_WORLD);
  
  for(currTask=0;currTask<NTasks;++currTask)
    {
      MPI_Barrier(MPI_COMM_WORLD);
      
      if(ThisTask == currTask)
	{
	  if(NumHalos > 0)
	    {
	      fp = fopen(name,"a");
	      
	      for(i=0;i<NumHalos;++i)
		fprintf(fp,formatString,Halos[i].index,
			Halos[i].radius,Halos[i].mass,Halos[i].vcirc,
			Halos[i].pos[0],Halos[i].pos[1],Halos[i].pos[2],
			Halos[i].vel[0],Halos[i].vel[1],Halos[i].vel[2]);
	      
	      fclose(fp);
	    }
	}
      
      MPI_Barrier(MPI_COMM_WORLD);
    }
}
#endif

void writepeaks(void)
{
  char name[MAX_FILENAME];
  sprintf(name,"%s",haloFinderData.PeakOutputFile);
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"writing peaks to '%s'\n",name);
#endif
  
#ifdef HAVEFITSIO
  writepeakshalos_name(name);
#else
  writepeakshalos_ascii(name);
#endif
}

void writehalos(void)
{
  char name[MAX_FILENAME];
  sprintf(name,"%s",haloFinderData.HaloOutputFile);
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"writing halos to '%s'\n",name);
#endif
  
#ifdef HAVEFITSIO
  writepeakshalos_name(name);
#else
  writepeakshalos_ascii(name);
#endif
}

void readhalos(void)
{
  char name[MAX_FILENAME];
  sprintf(name,"%s",haloFinderData.HaloOutputFile);
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"reading halos from '%s'\n",name);
#endif

  readpeakshalos(name);
}

void readpeaks(void)
{
  char name[MAX_FILENAME];
  
  sprintf(name,"%s",haloFinderData.PeakInputFile);

#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"reading peaks from '%s'\n",name);
#endif
  
  readpeakshalos(name);
}

static void readpeakshalos(char name[])
{
#ifndef HAVEFITSIO
  FILE *fp;
#endif
  long i,j;
  long TotNumHalos;
  long BaseHaloChunkSize,HaloChunkSize;
  Halo *tmpHalos,*chunkHalos;
  long *rank,chunk,NumChunks;
  int recvcount,*displs,*sendcounts;
  int coords[3],haloDestTask;
  size_t *sindex;
  float dx,dy,dz;
  Halo sourceHalo,saveHalo;
  long rankSource,rankSave,dest;
  
  dx = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[0]);
  dy = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[1]);
  dz = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[2]);
  
  if(ThisTask == 0)
    {
#ifdef HAVEFITSIO
      TotNumHalos = read_totnumpeakshalos_fits(name);
#else
      fp = fopen(name,"r");
      assert(fp != NULL);
      TotNumHalos = fnumlines(fp);
#endif
    }
  MPI_Bcast(&TotNumHalos,1,MPI_LONG,0,MPI_COMM_WORLD);
  
#ifdef HAVEFITSIO
  BaseHaloChunkSize = read_chunksizepeakshalos_fits(name);
#else
  BaseHaloChunkSize = (long) (IO_CHUNK_SIZE_MB*1024.0*1024.0/((double) (sizeof(Halo))));
#endif
  NumChunks = TotNumHalos/BaseHaloChunkSize;
  if(NumChunks*BaseHaloChunkSize != TotNumHalos)
    ++NumChunks;
  
  if(ThisTask == 0)
    {
      chunkHalos = (Halo*)malloc(sizeof(Halo)*BaseHaloChunkSize);
      assert(chunkHalos != NULL);
      sendcounts = (int*)malloc(sizeof(int)*NTasks);
      assert(sendcounts != NULL);
      displs = (int*)malloc(sizeof(int)*NTasks);
      assert(displs != NULL);
    }
  
  for(chunk=0;chunk<NumChunks;++chunk)
    {
      if(ThisTask == 0)
	{
	  HaloChunkSize = BaseHaloChunkSize;
	  if(chunk*BaseHaloChunkSize + HaloChunkSize > TotNumHalos)
	    HaloChunkSize = TotNumHalos - chunk*BaseHaloChunkSize;
	  
#ifdef HAVEFITSIO
	  read_peakshaloschunk_fits(name,chunkHalos,BaseHaloChunkSize,HaloChunkSize,chunk);
#else 
	  for(i=0;i<HaloChunkSize;++i)
	    fscanf(fp,formatStringRead,&(chunkHalos[i].index),
		   &(chunkHalos[i].radius),&(chunkHalos[i].mass),&(chunkHalos[i].vcirc),
		   &(chunkHalos[i].pos[0]),&(chunkHalos[i].pos[1]),&(chunkHalos[i].pos[2]),
		   &(chunkHalos[i].vel[0]),&(chunkHalos[i].vel[1]),&(chunkHalos[i].vel[2]));
	  
#endif  
	  rank = (long*)malloc(sizeof(long)*HaloChunkSize);
	  assert(rank != NULL);
	  for(i=0;i<HaloChunkSize;++i)
	    {
	      coords[0] = chunkHalos[i].pos[0]/dx;
	      coords[1] = chunkHalos[i].pos[1]/dy;
	      coords[2] = chunkHalos[i].pos[2]/dz;
	      for(j=0;j<3;++j)
		{
#ifndef LIGHTCONE
		  while(coords[j] < 0)
		    coords[j] += NumDomains[j];
		  while(coords[j] >= NumDomains[j])
		    coords[j] -= NumDomains[j];
#else
		  if(coords[j] < 0)
		    coords[j] = 0;
		  if(coords[j] >= NumDomains[j])
		    coords[j] = NumDomains[j]-1;
#endif
		}

	      MPI_Cart_rank(cartComm,coords,&haloDestTask);
	      rank[i] = (long) (haloDestTask);
	    }
	  
	  //set up for scatter
	  for(i=0;i<NTasks;++i)
	    sendcounts[i] = 0;
	  for(i=0;i<HaloChunkSize;++i)
	    sendcounts[rank[i]] += 1;
	  displs[0] = 0;
	  for(i=1;i<NTasks;++i)
	    displs[i] = displs[i-1]+sendcounts[i-1];
	  
	  //sort halos
	  sindex = (size_t*)malloc(sizeof(size_t)*HaloChunkSize);
	  assert(sindex != NULL);
	  gsl_sort_long_index(sindex,rank,(size_t) 1,(size_t) HaloChunkSize);
	  
	  for(i=0;i<HaloChunkSize;++i)
	    rank[sindex[i]] = i;
	  free(sindex);
	  
	  for(i=0;i<HaloChunkSize;++i) /* reoder with an in-place algorithm - see Gadget-2 for details - destroys rank */
	    {
	      if(i != rank[i])
		{
		  sourceHalo = chunkHalos[i];
		  rankSource = rank[i];
		  dest = rank[i];
		  
		  do
		    {
		      saveHalo = chunkHalos[dest];
		      rankSave = rank[dest];
		      
		      chunkHalos[dest] = sourceHalo;
		      rank[dest] = rankSource;
		      
		      if(dest == i)
			break;
		      
		      sourceHalo = saveHalo;
		      rankSource = rankSave;
		      
		      dest = rankSource;
		    }
		  while(1);
		}
	    }
	  
	  free(rank);
	}
      
      //send each task the number of halos it will recv
      MPI_Scatter(sendcounts,1,MPI_INT,&recvcount,1,MPI_INT,0,MPI_COMM_WORLD);
      if(recvcount > 0)
	{
	  tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(NumHalos+recvcount));
	  assert(tmpHalos != NULL);
	  Halos = tmpHalos;
	}

      //send the Halos
      MPI_Scatterv(chunkHalos,sendcounts,displs,MPI_Halo,
		   Halos+NumHalos,recvcount,MPI_Halo,
		   0,MPI_COMM_WORLD);
      NumHalos += recvcount;
    }
      
  if(ThisTask == 0)
    {
#ifndef HAVEFITSIO
      fclose(fp);
#endif
      free(chunkHalos);
      free(sendcounts);
      free(displs);
    }
}
